import java.io.BufferedWriter;
import java.io.FileWriter;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import org.eclipse.jdt.core.dom.ASTNode;
import org.eclipse.jdt.core.dom.Block;
import org.eclipse.jdt.core.dom.ForStatement;
import org.eclipse.jdt.core.dom.IfStatement;
import org.eclipse.jdt.core.dom.MethodDeclaration;
import org.eclipse.jdt.core.dom.Statement;
import org.eclipse.jdt.core.dom.WhileStatement;


public class CFG {
	private CFGNode entry;
	private CFGNode exit;
	private MethodDeclaration methodSig;
	private ArrayList<CFGNode> breaks;
	
	public static CFG buildCFG(MethodDeclaration repOk) {
		CFG temp = new CFG();
		temp.createCFG(repOk);
		return temp;
	}

	@SuppressWarnings("unchecked")
	private void createCFG(MethodDeclaration repOk) {
		methodSig = repOk;
		entry = new CFGNode(CFGNode.CFG_ENTRY);
		exit = new CFGNode(CFGNode.CFG_EXIT);
		breaks = new ArrayList<CFGNode>();
		CFGNode current = entry;
		Block body = repOk.getBody();
		List<Statement> statements = body.statements();
		
		for (Statement st : statements) {
			if (st.getNodeType() == ASTNode.RETURN_STATEMENT)
			{
				current.addStatement(st);
				current.addChild(exit);
				break;
			}
			else
				current = addStatement(st,current);
		}
		
		if(current != exit)
			current.addChild(exit);
		
		System.out.println("CFG is created");				
	}

	private CFGNode addStatement(Statement st, CFGNode currentNode) {
		if (st.getNodeType() == ASTNode.IF_STATEMENT)
			return addIf(st, currentNode);
		
		else if (st.getNodeType() == ASTNode.WHILE_STATEMENT)
			return addWhile((WhileStatement)st, currentNode);
		
		else if (st.getNodeType() == ASTNode.FOR_STATEMENT)
			return addForLoop((ForStatement)st, currentNode);
		
		else if (st.getNodeType() == ASTNode.BLOCK)
			for (Object line : ((Block)st).statements()) 
			{
				currentNode = addStatement((Statement)line,currentNode);
			}
		
		else
			currentNode.addStatement(st);		
		
		return currentNode;
	}

	@SuppressWarnings("unchecked")
	private CFGNode addForLoop(ForStatement st, CFGNode currentNode) {
		Block body = (Block) st.getBody();
		List<Statement> statements = body.statements();
		
		CFGNode forStatement = new CFGNode();
		CFGNode forExit = new CFGNode();
		CFGNode forBody = new CFGNode();
		
		currentNode.addChild(forStatement);
		currentNode.setNext(forStatement);
		forStatement.addStatement(st);
		forStatement.addChild(forBody);
		forStatement.addChild(forExit);
		
		currentNode = forBody;
		
		for (Statement s : statements) {
			if (s.getNodeType() == ASTNode.BREAK_STATEMENT){
				currentNode.addStatement(s);
				breaks.add(currentNode);
				break;
			}
			else if (s.getNodeType() == ASTNode.RETURN_STATEMENT){
				currentNode.addStatement(s);
				currentNode.addChild(this.exit);
				break;
			}
			else if (s.getNodeType() == ASTNode.CONTINUE_STATEMENT){
				currentNode.addStatement(s);
				currentNode.addChild(forStatement);
			}
			else
				currentNode = addStatement(s, currentNode);
		}
		
		currentNode.addChild(forStatement);
		
		for (CFGNode n : breaks) 
			n.addChild(forExit);
		breaks.clear();
		forStatement.setNext(forExit);
		
		return forExit;
		
	}

	@SuppressWarnings("unchecked")
	private CFGNode addWhile(WhileStatement whileSt, CFGNode currentNode) {
		CFGNode whileEntry = new CFGNode();
		CFGNode whileTrue = new CFGNode();
		CFGNode whileFalse = new CFGNode();
		
		currentNode.addChild(whileEntry);
		currentNode.setNext(whileEntry);
		whileEntry.addChild(whileTrue);
		whileEntry.addChild(whileFalse);
		whileEntry.addStatement(whileSt);
		
		currentNode = whileTrue;
		
		Block body = (Block) whileSt.getBody();
		List<Statement> statements = body.statements();

		for (Statement st : statements) {
			if (st.getNodeType() == ASTNode.BREAK_STATEMENT){
				currentNode.addStatement(st);
				breaks.add(currentNode);
				break;
			}
			else if (st.getNodeType() == ASTNode.RETURN_STATEMENT){
				currentNode.addStatement(st);
				currentNode.addChild(this.exit);
				break;
			}
			else if (st.getNodeType() == ASTNode.CONTINUE_STATEMENT){
				currentNode.addStatement(st);
				currentNode.addChild(whileEntry);
			}
			else
				currentNode = addStatement(st, currentNode);
		}	
		
		currentNode.addChild(whileEntry);
		for (CFGNode n : breaks) 
			n.addChild(whileFalse);
		
		breaks.clear();
		whileEntry.setNext(whileFalse);
		return whileFalse;
	}


	private CFGNode addIf(Statement st, CFGNode conditionNode) {
		
		conditionNode.addStatement(st);
		
		IfStatement ifSt = (IfStatement) st;
		Statement thenSt = ifSt.getThenStatement();
		Statement elseSt = ifSt.getElseStatement();
		
		CFGNode endNode = new CFGNode();
		
		CFGNode c1 = new CFGNode();
		conditionNode.addChild(c1);
		c1 = addStatement(thenSt, c1);
		if (c1.hasBreak())
			breaks.add(c1);
		else if (c1.hasReturn())
			c1.addChild(this.exit);
		else
			c1.addChild(endNode);
		
		if (elseSt!=null){
			CFGNode c2 = new CFGNode();
			conditionNode.addChild(c2);
			c2 = addStatement(elseSt, c2);
			if (c2.hasBreak())
				breaks.add(c2);
			else if (c2.hasReturn())
				c2.addChild(this.exit);
			else
				c2.addChild(endNode);
		}else
			conditionNode.addChild(endNode);
		
		conditionNode.setNext(endNode);
		return endNode;
	}

	public void writeToFile(String fileName) throws IOException {
		BufferedWriter writer = new BufferedWriter(new FileWriter(fileName));
		writer.write(methodSig.modifiers().get(0).toString() + " ");
		writer.write(methodSig.getReturnType2().toString() + " " + methodSig.getName().toString() + "(){");
		writer.newLine();
		entry.print(writer,"\t");
		writer.write("}");
		writer.newLine();
		writer.close();
	}
	
	public CFGNode getEntry() {
		return entry;
	}

	public void setEntry(CFGNode entry) {
		this.entry = entry;
	}

	public CFGNode getExit() {
		return exit;
	}

	public void setExit(CFGNode exit) {
		this.exit = exit;
	}
	
}
